import numpy as np
from collections import defaultdict


class NaiveBayes:
    def __init__(self):
        self.class_prob = {}
        self.feature_prob = {}
        self.classes = []

    def fit(self, X, y):
        """训练模型"""
        n_samples, n_features = len(X), len(X[0])
        self.classes = list(set(y))

        # 计算每个类别的先验概率
        for c in self.classes:
            self.class_prob[c] = sum(1 for label in y if label == c) / n_samples

        # 计算每个特征在每个类别下的条件概率
        self.feature_prob = defaultdict(lambda: defaultdict(dict))

        for c in self.classes:
            # 获取当前类别的所有样本
            class_samples = [X[i] for i in range(n_samples) if y[i] == c]

            for feature_idx in range(n_features):
                feature_values = [sample[feature_idx] for sample in class_samples]

                # 计算每个特征值的概率（使用拉普拉斯平滑）
                unique_values = list(set([sample[feature_idx] for sample in X]))
                total_count = len(feature_values)

                for value in unique_values:
                    count = sum(1 for v in feature_values if v == value)
                    self.feature_prob[c][feature_idx][value] = (count + 1) / (total_count + len(unique_values))

    def predict(self, X):
        """预测"""
        predictions = []
        for sample in X:
            max_prob = -1
            best_class = None

            for c in self.classes:
                # 计算后验概率 P(class|features) ∝ P(class) * ∏P(feature|class)
                prob = self.class_prob[c]

                for feature_idx, value in enumerate(sample):
                    if value in self.feature_prob[c][feature_idx]:
                        prob *= self.feature_prob[c][feature_idx][value]
                    else:
                        # 如果特征值没出现过，使用很小的概率
                        prob *= 1e-6

                if prob > max_prob:
                    max_prob = prob
                    best_class = c

            predictions.append(best_class)

        return predictions


# 测试代码
if __name__ == "__main__":
    # 简单的天气数据集
    X_train = [
        ['Sunny', 'Hot', 'High', 'Weak'],
        ['Sunny', 'Hot', 'High', 'Strong'],
        ['Overcast', 'Hot', 'High', 'Weak'],
        ['Rain', 'Mild', 'High', 'Weak'],
        ['Rain', 'Cool', 'Normal', 'Weak'],
        ['Rain', 'Cool', 'Normal', 'Strong'],
        ['Overcast', 'Cool', 'Normal', 'Strong'],
        ['Sunny', 'Mild', 'High', 'Weak'],
        ['Sunny', 'Cool', 'Normal', 'Weak'],
        ['Rain', 'Mild', 'Normal', 'Weak']
    ]

    y_train = ['No', 'No', 'Yes', 'Yes', 'Yes', 'No', 'Yes', 'No', 'Yes', 'Yes']

    X_test = [['Sunny', 'Cool', 'High', 'Strong']]

    nb = NaiveBayes()
    nb.fit(X_train, y_train)
    predictions = nb.predict(X_test)

    print("预测结果:", predictions)
    print("实际应该是: No")